"""
任务分区处理器

这个模块实现了任务分区功能，能够将可分区的任务分解为多个并行子任务，
以优化资源利用和项目执行效率。
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
import math
import logging

from .simplified_data_models import Activity, Precedence, ProjectData, SubdivisionConfig
from .exceptions import SchedulerError


logger = logging.getLogger(__name__)


class SubdivisionError(SchedulerError):
    """分区处理错误"""
    
    def __init__(self, message: str, activity_id: Optional[int] = None, subdivision_count: Optional[int] = None):
        super().__init__(message, "SUBDIVISION_ERROR")
        self.activity_id = activity_id
        self.subdivision_count = subdivision_count


@dataclass
class SubdividedActivity:
    """分区后的子任务数据模型"""
    
    id: int
    name: str
    duration: int
    parent_id: int
    subdivision_index: int
    original_duration: int
    worker_requirements: Dict[str, int] = field(default_factory=dict)
    can_subdivide: bool = False
    subdivided: bool = True
    
    def to_activity(self) -> Activity:
        """转换为标准Activity对象"""
        activity = Activity(
            id=self.id,
            name=self.name,
            duration=self.duration,
            can_subdivide=self.can_subdivide,
            original_duration=self.original_duration,
            subdivided=self.subdivided,
            worker_requirements=self.worker_requirements.copy()
        )
        return activity


class SubdivisionProcessor:
    """任务分区处理器
    
    负责处理任务的分区逻辑，包括：
    - 识别可分区任务
    - 创建子任务
    - 分配工期和资源
    - 处理依赖关系
    """
    
    def __init__(self):
        self.next_activity_id = 10000  # 子任务ID起始值，避免与原任务冲突
        
    def subdivide_activities(
        self, 
        project_data: ProjectData, 
        subdivision_count: int
    ) -> ProjectData:
        """将可分区任务分解为子任务
        
        Args:
            project_data: 原始项目数据
            subdivision_count: 分区数量
            
        Returns:
            包含分区后任务的新项目数据
            
        Raises:
            SubdivisionError: 分区处理失败时
        """
        if subdivision_count < 2:
            raise SubdivisionError(
                f"分区数量必须大于等于2，当前值: {subdivision_count}",
                subdivision_count=subdivision_count
            )
            
        logger.info(f"开始任务分区处理，分区数量: {subdivision_count}")
        
        try:
            # 创建新的项目数据副本
            new_activities = []
            subdivision_mapping = {}  # 原任务ID -> 子任务ID列表的映射
            
            # 处理每个任务
            for activity in project_data.activities:
                if activity.can_subdivide:
                    # 创建子任务
                    sub_activities = self._create_sub_activities(activity, subdivision_count)
                    new_activities.extend([sub_act.to_activity() for sub_act in sub_activities])
                    subdivision_mapping[activity.id] = [sub_act.id for sub_act in sub_activities]
                    
                    logger.info(f"任务 {activity.id}({activity.name}) 已分区为 {len(sub_activities)} 个子任务")
                else:
                    # 保持原任务不变
                    new_activities.append(activity)
            
            # 处理前置关系
            new_precedences = self._update_precedences(
                project_data.precedences, 
                subdivision_mapping
            )
            
            # 创建新的项目数据
            new_project_data = ProjectData(
                activities=new_activities,
                precedences=new_precedences,
                resources=project_data.resources.copy(),
                resource_demands=project_data.resource_demands.copy()
            )
            
            logger.info(f"分区处理完成，原任务数: {len(project_data.activities)}, 新任务数: {len(new_activities)}")
            return new_project_data
            
        except Exception as e:
            logger.error(f"任务分区处理失败: {e}")
            raise SubdivisionError(f"任务分区处理失败: {e}")
    
    def _create_sub_activities(
        self, 
        original_activity: Activity, 
        subdivision_count: int
    ) -> List[SubdividedActivity]:
        """创建子任务
        
        Args:
            original_activity: 原始任务
            subdivision_count: 分区数量
            
        Returns:
            子任务列表
        """
        sub_activities = []
        
        # 计算每个子任务的基础工期
        base_duration = original_activity.duration // subdivision_count
        remainder = original_activity.duration % subdivision_count
        
        logger.debug(f"分区任务 {original_activity.name}: 总工期={original_activity.duration}, "
                    f"基础工期={base_duration}, 余数={remainder}")
        
        # 创建子任务
        for i in range(subdivision_count):
            # 分配工期（余数分配给前面的子任务）
            duration = base_duration + (1 if i < remainder else 0)
            
            # 确保每个子任务至少有1天工期
            if duration < 1:
                duration = 1
            
            # 生成子任务ID和名称
            sub_activity_id = self._generate_sub_activity_id()
            sub_activity_name = f"{original_activity.name} 分区{i+1}"
            
            # 分配工种需求
            sub_worker_requirements = self._distribute_worker_requirements(
                original_activity.worker_requirements, 
                subdivision_count, 
                i
            )
            
            # 创建子任务
            sub_activity = SubdividedActivity(
                id=sub_activity_id,
                name=sub_activity_name,
                duration=duration,
                parent_id=original_activity.id,
                subdivision_index=i + 1,
                original_duration=original_activity.duration,
                worker_requirements=sub_worker_requirements,
                can_subdivide=False,  # 子任务不能再次分区
                subdivided=True
            )
            
            sub_activities.append(sub_activity)
            
            logger.debug(f"创建子任务: {sub_activity_name} (ID: {sub_activity_id}, 工期: {duration})")
        
        return sub_activities
    
    def _distribute_worker_requirements(
        self, 
        original_requirements: Dict[str, int], 
        subdivision_count: int, 
        subdivision_index: int
    ) -> Dict[str, int]:
        """分配工种需求到子任务
        
        Args:
            original_requirements: 原始工种需求
            subdivision_count: 分区数量
            subdivision_index: 当前子任务索引（从0开始）
            
        Returns:
            子任务的工种需求
        """
        if not original_requirements:
            return {}
        
        distributed_requirements = {}
        
        for worker_type, total_count in original_requirements.items():
            # 计算每个子任务的基础工人数量
            base_count = total_count // subdivision_count
            remainder = total_count % subdivision_count
            
            # 分配工人数量（余数分配给前面的子任务）
            worker_count = base_count + (1 if subdivision_index < remainder else 0)
            
            # 确保每个子任务至少有工人（如果原任务有该工种需求）
            if worker_count == 0 and total_count > 0:
                worker_count = 1
            
            if worker_count > 0:
                distributed_requirements[worker_type] = worker_count
                
            logger.debug(f"工种 {worker_type}: 总数={total_count}, 分区{subdivision_index+1}分配={worker_count}")
        
        return distributed_requirements
    
    def _update_precedences(
        self, 
        original_precedences: List[Precedence], 
        subdivision_mapping: Dict[int, List[int]]
    ) -> List[Precedence]:
        """更新前置关系以适应分区后的任务
        
        Args:
            original_precedences: 原始前置关系列表
            subdivision_mapping: 任务分区映射（原任务ID -> 子任务ID列表）
            
        Returns:
            更新后的前置关系列表
        """
        new_precedences = []
        
        for precedence in original_precedences:
            from_activity = precedence.from_activity
            to_activity = precedence.to_activity
            
            # 获取前置任务和后续任务的子任务列表
            from_sub_activities = subdivision_mapping.get(from_activity, [from_activity])
            to_sub_activities = subdivision_mapping.get(to_activity, [to_activity])
            
            # 根据依赖类型创建新的前置关系
            if precedence.dependency_type == "FS":  # Finish-to-Start
                # 前置任务的最后一个子任务 -> 后续任务的第一个子任务
                last_from = from_sub_activities[-1] if isinstance(from_sub_activities, list) else from_sub_activities
                first_to = to_sub_activities[0] if isinstance(to_sub_activities, list) else to_sub_activities
                
                new_precedence = Precedence(
                    from_activity=last_from,
                    to_activity=first_to,
                    dependency_type=precedence.dependency_type,
                    lag=precedence.lag
                )
                new_precedences.append(new_precedence)
                
            elif precedence.dependency_type == "SS":  # Start-to-Start
                # 前置任务的第一个子任务 -> 后续任务的第一个子任务
                first_from = from_sub_activities[0] if isinstance(from_sub_activities, list) else from_sub_activities
                first_to = to_sub_activities[0] if isinstance(to_sub_activities, list) else to_sub_activities
                
                new_precedence = Precedence(
                    from_activity=first_from,
                    to_activity=first_to,
                    dependency_type=precedence.dependency_type,
                    lag=precedence.lag
                )
                new_precedences.append(new_precedence)
                
            elif precedence.dependency_type == "FF":  # Finish-to-Finish
                # 前置任务的最后一个子任务 -> 后续任务的最后一个子任务
                last_from = from_sub_activities[-1] if isinstance(from_sub_activities, list) else from_sub_activities
                last_to = to_sub_activities[-1] if isinstance(to_sub_activities, list) else to_sub_activities
                
                new_precedence = Precedence(
                    from_activity=last_from,
                    to_activity=last_to,
                    dependency_type=precedence.dependency_type,
                    lag=precedence.lag
                )
                new_precedences.append(new_precedence)
                
            elif precedence.dependency_type == "SF":  # Start-to-Finish
                # 前置任务的第一个子任务 -> 后续任务的最后一个子任务
                first_from = from_sub_activities[0] if isinstance(from_sub_activities, list) else from_sub_activities
                last_to = to_sub_activities[-1] if isinstance(to_sub_activities, list) else to_sub_activities
                
                new_precedence = Precedence(
                    from_activity=first_from,
                    to_activity=last_to,
                    dependency_type=precedence.dependency_type,
                    lag=precedence.lag
                )
                new_precedences.append(new_precedence)
            
            logger.debug(f"更新前置关系: {from_activity} -> {to_activity} "
                        f"({precedence.dependency_type}) 转换为新关系")
        
        return new_precedences
    
    def _generate_sub_activity_id(self) -> int:
        """生成子任务ID"""
        current_id = self.next_activity_id
        self.next_activity_id += 1
        return current_id
    
    def get_subdivision_summary(self, project_data: ProjectData) -> Dict[str, int]:
        """获取分区处理摘要信息
        
        Args:
            project_data: 项目数据
            
        Returns:
            包含分区统计信息的字典
        """
        total_activities = len(project_data.activities)
        subdividable_activities = len([act for act in project_data.activities if act.can_subdivide])
        subdivided_activities = len([act for act in project_data.activities if act.subdivided])
        
        return {
            "total_activities": total_activities,
            "subdividable_activities": subdividable_activities,
            "subdivided_activities": subdivided_activities,
            "subdivision_ratio": subdivided_activities / total_activities if total_activities > 0 else 0
        }
    
    def validate_subdivision_result(self, original_data: ProjectData, subdivided_data: ProjectData) -> bool:
        """验证分区结果的正确性
        
        Args:
            original_data: 原始项目数据
            subdivided_data: 分区后项目数据
            
        Returns:
            验证是否通过
            
        Raises:
            SubdivisionError: 验证失败时
        """
        try:
            # 验证任务数量
            original_subdividable = len([act for act in original_data.activities if act.can_subdivide])
            original_non_subdividable = len([act for act in original_data.activities if not act.can_subdivide])
            
            subdivided_count = len([act for act in subdivided_data.activities if act.subdivided])
            non_subdivided_count = len([act for act in subdivided_data.activities if not act.subdivided])
            
            # 非分区任务数量应该保持不变
            if non_subdivided_count != original_non_subdividable:
                raise SubdivisionError(
                    f"非分区任务数量不匹配: 原始={original_non_subdividable}, 处理后={non_subdivided_count}"
                )
            
            # 验证工期分配
            for activity in subdivided_data.activities:
                if activity.subdivided and activity.duration <= 0:
                    raise SubdivisionError(
                        f"子任务 {activity.name} 的工期无效: {activity.duration}",
                        activity_id=activity.id
                    )
            
            logger.info("分区结果验证通过")
            return True
            
        except Exception as e:
            logger.error(f"分区结果验证失败: {e}")
            raise SubdivisionError(f"分区结果验证失败: {e}")


def create_subdivision_processor() -> SubdivisionProcessor:
    """创建分区处理器实例"""
    return SubdivisionProcessor()